import pandas as pd
from openpyxl import Workbook
from openpyxl.styles import PatternFill
from datetime import datetime, timedelta
import openpyxl
import os

class DataVisualizer:
    def __init__(self, save_path='./visualize_未归一化.xlsx', visual_days=5):
        self.ws = None
        self.save_path = save_path
        self.visual_days = visual_days
        if os.path.exists(self.save_path):
            self.wb = openpyxl.load_workbook(self.save_path)
        else:
            self.wb = Workbook()


    def visualize_data(self, selected_data: dict, all_data: dict, stock_code_name: dict, detect_date: str) -> None:
        """
        input: dict(),'stock_code': 'interval':[(start, end),...]
        :return:
        """
        # 根据分数排序selected_data
        selected_data = {key: value for key, value in
                         sorted(selected_data.items(), key=lambda x: x[1]['score'], reverse=True)}

        # 获取默认的工作表
        sheet_name = detect_date
        self.wb.create_sheet(sheet_name)
        # 写入表头
        self.ws = self.wb[sheet_name]
        self.ws.cell(row=1, column=1, value='ts_code')
        self.ws.cell(row=1, column=2, value='名称')
        self.ws.cell(row=1, column=3, value='得分')

        row_index = 2
        for code in selected_data.keys():
            col_index = 4
            self.ws.cell(row=row_index, column=1, value=code)
            self.ws.cell(row=row_index, column=2, value=all_data[code][detect_date]['name'])
            self.ws.cell(row=row_index, column=3, value=selected_data[code]['score'])
            for start, end in selected_data[code]['interval']:
                # 可视化每个区间之后的N天数据
                self.ws.cell(row=row_index - 1, column=col_index, value=start)
                self.ws.cell(row=row_index, column=col_index, value=all_data[code][start]['pct_chg'])
                col_index = col_index + 1
                if end in all_data[code]:
                    for i in range(self.visual_days):
                        end_i = self.add_days(end, i)
                        if end_i in all_data[code]:
                            self.ws.cell(row=row_index - 1, column=col_index, value=end_i)
                            self.ws.cell(row=row_index, column=col_index, value=all_data[code][end_i]['pct_chg'])
                            # 填充
                            fill = PatternFill(start_color=self.cell_color(all_data[code][end_i]['pct_chg']),
                                               end_color=self.cell_color(all_data[code][end_i]['pct_chg']),
                                               fill_type='solid')
                            # 将填充应用于单元格
                            self.ws.cell(row=row_index, column=col_index).fill = fill
                            col_index = col_index + 1
                col_index = col_index + 1
            row_index = row_index + 3

        self.wb.save(self.save_path)

    def add_days(self, start_date, n):
        # 将字符串日期转换为datetime对象
        date_obj = datetime.strptime(start_date, '%Y%m%d')
        # 计算n天后的日期
        new_date = date_obj + timedelta(days=n)
        # 将结果转换为字符串形式
        new_date_str = new_date.strftime('%Y%m%d')
        return new_date_str

    def cell_color(self, data):
        if isinstance(data, str):
            pass
        else:
            if 0.0 <= data < 5.0:
                return 'FF6666'  # rgb(250,128,114)' # 橙红色
            elif 5.0 <= data < 10.0:
                return 'FF0033'  # rgb(255,0,0)' # 红色
            elif data >= 10.0:
                return '990000'  # rgb(176,23,31)' # 印度红
            elif -5.0 <= data < 0.0:
                return '66FF66'  # rgb(127,255,0)'#黄绿色
            elif -10.0 <= data < -5.0:
                return '339933'  # rgb(0,201,87)'  # 翠绿色
            elif data < -10.0:
                return '006600'  # rgb(61,145,64)' # 钴绿色
            else:
                return 'FFFFFF'  # rgb(0,0,0)' # 白色
